package com.hzw.saas.web.shell.util;

import cn.hutool.core.io.IORuntimeException;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.ChannelType;
import cn.hutool.extra.ssh.JschRuntimeException;
import cn.hutool.extra.ssh.JschUtil;
import cn.hutool.extra.ssh.Sftp;
import com.hzw.saas.common.config.constants.RedisConsts;
import com.hzw.saas.common.config.util.RedisUtil;
import com.hzw.saas.web.shell.event.DeleteSessionEvent;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.Charset;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Jsch工具补充
 *
 * @author zzl
 * @see cn.hutool.extra.ssh.JschUtil
 * @since 06/07/2021
 */
@Slf4j
public class JschEnhanceUtil {

    /**
     * 重新尝试次数
     */
    public static final int RETRY_TIMES = 3;
    /**
     * 重新尝试间隔，ms
     */
    public static final long RETRY_INTERVAL = 100;

    /**
     * session持续时间，每次获取session将会刷新时间
     */
    private static final long SESSION_EXPIRE_TIME = 60 * 30;

    /**
     * 同步锁
     */
    private static final Map<String, Lock> keyLockMap = new ConcurrentHashMap<>();

    /**
     * 获得一个SSH会话，重用已经使用的会话
     * 此会话在6小时后会被回收，重新调用该方法会刷新回收时间
     *
     * @param host     主机
     * @param port     端口
     * @param username 用户名
     * @param password 密码
     * @param sid      sessionId, 确保唯一性
     * @return session
     */
    public static Session getSession(String host, int port, String username, String password, String sid) {
        // 加锁
        final Lock keyLock = keyLockMap.computeIfAbsent(sid, k -> new ReentrantLock());
        keyLock.lock();
        try {
            Session session = null;
            getSession:
            for (int i = 1; i <= RETRY_TIMES && (Objects.isNull(session) || !session.isConnected()); i++) {
                log.debug("尝试第[{}]次获取session,{sid: {}, username: {}, port: {}}", i, sid, username, port);
                session = JschEnhanceSessionPool.INSTANCE.getSession(host, port, username, password, sid);
                // 获取到的session未连接时,尝试连接
                for (int j = 1; j <= RETRY_TIMES && (Objects.nonNull(session) && !session.isConnected()); j++) {
                    try {
                        log.debug("尝试第[{}]次连接session,{sid: {}, username: {}, port: {}}", j, sid, username, port);
                        session.connect();
                        ThreadUtil.sleep(RETRY_INTERVAL);
                    } catch (JSchException e) {
                        JschEnhanceSessionPool.INSTANCE.close(sid, 1);
                        continue getSession;
                    }
                }
            }
            if (Objects.isNull(session) || !session.isConnected()) {
                log.error("获取session失败,{sid: {}, username: {}, port: {}}", sid, username, port);
                throw new RuntimeException("获取连接失败，请稍后尝试");
            } else {
                log.debug("获取session完成,{sid: {}, username: {}, port: {}}", sid, username, port);
            }
            // 设置redis过期监听
//            setExpireListen(sid);
            return session;
        } finally {
            // 解锁
            keyLock.unlock();
            keyLockMap.remove(sid);
        }
    }

    /**
     * 同JschUtil.exec(Session session, String cmd, Charset charset, OutputStream errStream)
     * 同时开启了伪终端(未开启时，无法获取hpc的错误信息)
     *
     * @see JschUtil
     */
    public static String exec(Session session, String cmd, Charset charset, OutputStream errStream) {
        if (Objects.isNull(charset)) {
            charset = CharsetUtil.CHARSET_UTF_8;
        }
        final ChannelExec channel = (ChannelExec) JschUtil.createChannel(session, ChannelType.EXEC);
        channel.setPty(true);
        channel.setCommand(StrUtil.bytes(cmd, charset));
        channel.setInputStream(null);
        channel.setErrStream(errStream);
        InputStream in = null;
        try {
            log.debug("ssh exec: {}", cmd);
            channel.connect();
            in = channel.getInputStream();
            return IoUtil.read(in, charset);
        } catch (IOException e) {
            throw new IORuntimeException(e);
        } catch (JSchException e) {
            throw new JschRuntimeException(e);
        } finally {
            IoUtil.close(in);
            JschUtil.close(channel);
        }
    }

    /**
     * 根据SID关闭session会话并从缓存移除
     *
     * @param sid
     */
    public static void closeSession(String sid) {
        log.debug("close session({})...", sid);
        JschEnhanceSessionPool.INSTANCE.close(sid, RETRY_TIMES);
        removeExpireListen(sid);
    }

    public static Sftp getSftp(String host, int port, String username, String password, String cid) {
        Session session = getSession(host, port, username, password, cid);
        return JschUtil.createSftp(session);
    }

    public static ChannelSftp getChannelSftp(String host, int port, String username, String password, String cid) {
        return getSftp(host, port, username, password, cid).getClient();
    }

    public static void closeChannelInSftp(Sftp sftp) {
        if (Objects.nonNull(sftp))
            JschUtil.close(sftp.getClient());
    }

    private static void setExpireListen(String sid) {
        String expireTriggerKey = RedisConsts.GROUP_LISTEN_EXPIRE_TRIGGER + sid;
        String expireEventKey = RedisConsts.GROUP_LISTEN_EXPIRE_EVENT + sid;
        if (RedisUtil.hasKey(expireTriggerKey)) {
            RedisUtil.expire(expireTriggerKey, SESSION_EXPIRE_TIME);
        } else {
            RedisUtil.set(expireTriggerKey, sid, SESSION_EXPIRE_TIME);
        }
        if (!RedisUtil.hasKey(expireEventKey)) {
            RedisUtil.set(expireEventKey, new DeleteSessionEvent(sid));
        }
    }

    private static void removeExpireListen(String sid) {
        String expireKey = RedisConsts.GROUP_LISTEN_EXPIRE_TRIGGER + sid;
        String dataKey = RedisConsts.GROUP_LISTEN_EXPIRE_EVENT + sid;
        RedisUtil.del(expireKey, dataKey);
    }

}
